#include <fastbevdetect.h>


void Fastbevdetect::Init()
{
    //读取配置文件
    YAML::Node config = YAML::LoadFile(config_file_path_);

    //模型位置
    model=config["model_dir"].as<std::string>();
    std::string precision=config["precision"].as<std::string>();
    std::string sensor_data_dir=config["sensor_data_dir"].as<std::string>();
    const char* sensor_data=sensor_data_dir.c_str();

    // 从tensorrt文件中读取模型
    printf("Create model\n");
    core = Fastbevdetect::create_core(model, precision);
    core->print();
    core->set_timer(true);

    cudaStreamCreate(&stream);

    // 读取标定参数
    printf("Create data %s/valid_c_idx.tensor \n", sensor_data);

    //见鬼  这么写会读取数据报错
    // auto valid_c_idx = nv::Tensor::load(nv::format("%s/valid_c_idx.tensor", data), false);

    std::string valid_c_idx_path=sensor_data_dir+"valid_c_idx.tensor";
    std::string valid_x_path=sensor_data_dir+"x.tensor";
    std::string valid_y_path=sensor_data_dir+"y.tensor";
    auto valid_c_idx = nv::Tensor::load(valid_c_idx_path, false);
    auto valid_x = nv::Tensor::load(valid_x_path, false);
    auto valid_y = nv::Tensor::load(valid_y_path, false);

    // 数据流更新，将标定参数更新到模型中
    printf("update data %s/valid_c_idx.tensor \n", sensor_data);
    core->update(valid_c_idx.ptr<float>(), valid_x.ptr<int64_t>(), valid_y.ptr<int64_t>(), stream);
    printf("update data end \n");

    // 从ros中接受的消息
    std::string front_left_channel=config["front_left_channel"].as<std::string>();
    std::string front_channel=config["front_channel"].as<std::string>();
    std::string front_right_channel=config["front_right_channel"].as<std::string>();
    std::string back_left_channel=config["back_left_channel"].as<std::string>();
    std::string back_channel=config["back_channel"].as<std::string>();
    std::string back_right_channel=config["back_right_channel"].as<std::string>();
    std::string visual_channel=config["visual_channel"].as<std::string>();
    std::string image_publish_channel=config["image_publish_channel"].as<std::string>();
    visual_frame=config["visual_frame"].as<std::string>();

    // 初始化监听函数
    front_left_image_sub_ = new message_filters::Subscriber<sensor_msgs::Image>(
        nh_, front_left_channel, 10);
    front_image_sub_ = new message_filters::Subscriber<sensor_msgs::Image>(
        nh_, front_channel, 10);
    front_right_image_sub_ = new message_filters::Subscriber<sensor_msgs::Image>(
        nh_, front_right_channel, 10);
    back_left_image_sub_ = new message_filters::Subscriber<sensor_msgs::Image>(
        nh_, back_left_channel, 10);
    back_image_sub_ = new message_filters::Subscriber<sensor_msgs::Image>(
        nh_, back_channel, 10);
    back_right_image_sub_ = new message_filters::Subscriber<sensor_msgs::Image>(
        nh_, back_right_channel, 10);

    // 可视化发布节点
    visual_publisher = nh_.advertise<visualization_msgs::MarkerArray>(
      visual_channel, 1);

    image_publisher = nh_.advertise<sensor_msgs::Image>(
      image_publish_channel, 1);

    
    // 多消息同步
    message_filters::Synchronizer<MySyncPolicy> *sync;

    sync = new message_filters::Synchronizer<MySyncPolicy>(MySyncPolicy(10),
                                                            *front_left_image_sub_, *front_image_sub_, *front_right_image_sub_,
                                                            *back_left_image_sub_, *back_image_sub_, *back_right_image_sub_ );
    sync->registerCallback(boost::bind(&Fastbevdetect::HandleSensorMessage,
                                        this, _1, _2, _3, _4, _5, _6));
    ros::spin();
    ros::shutdown();
} 

void Fastbevdetect::ImageMsg2Mat(sensor_msgs::ImageConstPtr msg,cv::Mat &image)
{
    // 从sensor_msgs::Image转换成cv::Mat
    cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::RGB8);
    image=cv_ptr->image;
}

void Fastbevdetect::HandleSensorMessage(
    const sensor_msgs::ImageConstPtr &front_left_image,
    const sensor_msgs::ImageConstPtr &front_image,
    const sensor_msgs::ImageConstPtr &front_right_image,
    const sensor_msgs::ImageConstPtr &back_left_image,
    const sensor_msgs::ImageConstPtr &back_image,
    const sensor_msgs::ImageConstPtr &back_right_image
    )
{
    std::cout << "INTO CALLBACK ,START INFER " <<std::endl;
    //图像格式转换
    cv::Mat front_left_image_mat, front_image_mat, front_right_image_mat,
            back_left_image_mat, back_image_mat, back_right_image_mat;
    ImageMsg2Mat(front_left_image,front_left_image_mat);
    ImageMsg2Mat(front_image,front_image_mat);
    ImageMsg2Mat(front_right_image,front_right_image_mat);
    ImageMsg2Mat(back_left_image,back_left_image_mat);
    ImageMsg2Mat(back_image,back_image_mat);
    ImageMsg2Mat(back_right_image,back_right_image_mat);

    std::vector<unsigned char*> images;

    images.push_back(front_image_mat.data);
    images.push_back(front_right_image_mat.data);
    images.push_back(front_left_image_mat.data);
    
    images.push_back(back_image_mat.data);
    images.push_back(back_left_image_mat.data);
    images.push_back(back_right_image_mat.data);
    
    
    
    
    saveImagesToBin(images, "test.bin");
    // 推理
    auto bboxes = core->forward((const unsigned char**)images.data(), stream);
    std::cout << "DET OBJ NUM: " << bboxes.size() <<std::endl;

    // 将结果转换为ros可视化marker
    auto VisMarkerArray = std::make_shared<visualization_msgs::MarkerArray>();
    std::vector<visualization_msgs::Marker> markes=compute_points(bboxes);
    VisMarkerArray->markers = markes;

    //image publish
    std::vector<cv::Mat> images_visual{front_left_image_mat, front_image_mat, front_right_image_mat,
            back_left_image_mat, back_image_mat, back_right_image_mat};
    visual_image(bboxes,images_visual);
    cv::Mat combinedImage = combineImages(images_visual);
    cv_bridge::CvImage cv_image;
    cv_image.header.stamp = ros::Time::now();
    cv_image.header.frame_id = "camera_frame";
    cv_image.encoding = sensor_msgs::image_encodings::BGR8;
    cv_image.image = combinedImage;

    // 发布图像消息
    image_publisher.publish(cv_image.toImageMsg());
    // cv::imwrite("test.jpg",combinedImage);


    // 、、发送消息
    visual_publisher.publish(*VisMarkerArray);


    images.clear();
}

cv::Point2f Fastbevdetect::projectPoint(const Point3D& point3D, const std::vector<double>& intrinsic, const std::vector<double>& rotation, const std::vector<double>& translation) 
{
    // 构建旋转矩阵
    cv::Mat R(3, 3, CV_64F);
    R.at<double>(0, 0) = rotation[0];
    R.at<double>(0, 1) = rotation[1];
    R.at<double>(0, 2) = rotation[2];
    R.at<double>(1, 0) = rotation[3];
    R.at<double>(1, 1) = rotation[4];
    R.at<double>(1, 2) = rotation[5];
    R.at<double>(2, 0) = rotation[6];
    R.at<double>(2, 1) = rotation[7];
    R.at<double>(2, 2) = rotation[8];

    // 构建平移向量
    cv::Mat t(3, 1, CV_64F);
    t.at<double>(0, 0) = translation[0];
    t.at<double>(1, 0) = translation[1];
    t.at<double>(2, 0) = translation[2];

    // 构建内参矩阵
    cv::Mat K(3, 3, CV_64F);
    K.at<double>(0, 0) = intrinsic[0];
    K.at<double>(0, 1) = intrinsic[1];
    K.at<double>(0, 2) = intrinsic[2];
    K.at<double>(1, 0) = intrinsic[3];
    K.at<double>(1, 1) = intrinsic[4];
    K.at<double>(1, 2) = intrinsic[5];
    K.at<double>(2, 0) = intrinsic[6];
    K.at<double>(2, 1) = intrinsic[7];
    K.at<double>(2, 2) = intrinsic[8];

    // 将 3D 点转换为齐次坐标
    cv::Mat point3DH(4, 1, CV_64F);
    point3DH.at<double>(0, 0) = point3D.x;
    point3DH.at<double>(1, 0) = point3D.y;
    point3DH.at<double>(2, 0) = point3D.z;
    point3DH.at<double>(3, 0) = 1.0;

    // 构建变换矩阵
    cv::Mat RT(3, 4, CV_64F);
    cv::hconcat(R, t, RT);

    // 进行投影变换
    cv::Mat point2DH = K * RT * point3DH;

    // 转换为非齐次坐标
    cv::Point2f point2D;
    point2D.x = point2DH.at<double>(0, 0) / point2DH.at<double>(2, 0);
    point2D.y = point2DH.at<double>(1, 0) / point2DH.at<double>(2, 0);

    return point2D;
}


void Fastbevdetect::saveImagesToBin(const std::vector<unsigned char*>& images, const std::string& filename) {
    std::ofstream outFile(filename, std::ios::binary);
    if (!outFile) {
        std::cerr << "无法打开文件进行写入。" << std::endl;
        return;
    }

    // 先写入图像的数量
    size_t numImages = images.size();
    // outFile.write(reinterpret_cast<const char*>(&numImages), sizeof(numImages));

    for (size_t i = 0; i < numImages; ++i) {
        // 假设每个图像数据的长度固定为 1024 字节，这里可以根据实际情况修改
        const size_t imageSize = 1920*1080*3;
        outFile.write(reinterpret_cast<const char*>(images[i]), imageSize);
    }

    outFile.close();
}

cv::Mat Fastbevdetect::combineImages(const std::vector<cv::Mat>& images) 
{
    if (images.size() != 6) {
        std::cerr << "The number of input images must be 6." << std::endl;
        return cv::Mat();
    }

    // 获取图像的尺寸
    int imageWidth = images[0].cols;
    int imageHeight = images[0].rows;

    // 计算组合后大图的尺寸
    int combinedWidth = 3 * imageWidth;
    int combinedHeight = 2 * imageHeight;

    // 创建组合后的大图
    cv::Mat combinedImage(combinedHeight, combinedWidth, images[0].type());

    // 将每张图像复制到组合图的对应位置
    for (int i = 0; i < 2; ++i) {
        for (int j = 0; j < 3; ++j) {
            int index = i * 3 + j;
            cv::Rect roi(j * imageWidth, i * imageHeight, imageWidth, imageHeight);
            images[index].copyTo(combinedImage(roi));
        }
    }

    return combinedImage;
}

void Fastbevdetect::visual_image(std::vector<fastbev::post::transbbox::BoundingBox> bboxes,std::vector<cv::Mat> &images)
{
    for (int j=0;j<images.size();j++)
    {
        cv::Mat test_image=images[j];

        std::ifstream file(json_file[j]);
        json calib;
        file >> calib;

        // 关闭文件
        file.close();
        std::vector<double> intrinsic = calib["intrinsic"].get<std::vector<double>>();
        std::vector<double> distortion = calib["distortion"].get<std::vector<double>>();
        std::vector<double> rotation = calib["rotation"].get<std::vector<double>>();
        std::vector<double> translation = calib["translation"].get<std::vector<double>>();
        
        cv::Mat intrinsic_mat = (cv::Mat_<double>(3, 3) << intrinsic[0], intrinsic[1], intrinsic[2],
                                                intrinsic[3], intrinsic[4], intrinsic[5],
                                                intrinsic[6], intrinsic[7], intrinsic[8]);
        cv::Mat distort_mat = (cv::Mat_<double>(1, 4) << distortion[0], distortion[1], distortion[2], distortion[3]);
        cv::Mat map1, map2;
        cv::Size imageSize(test_image.cols, test_image.rows);
        cv::fisheye::initUndistortRectifyMap(intrinsic_mat, distort_mat, cv::Mat::eye(3, 3, CV_64F), intrinsic_mat, imageSize, CV_16SC2, map1, map2);

        // 应用去畸变映射
        cv::remap(test_image, test_image, map1, map2, cv::INTER_LINEAR, cv::BORDER_CONSTANT);

        for (int i=0;i<bboxes.size();i++)
        {
            fastbev::post::transbbox::BoundingBox box=bboxes[i];
            float x=box.position.x;
            float y=box.position.y;
            float z=box.position.z;
            z=z+1;
            float w=box.size.w;
            float l=box.size.l;
            float h=box.size.h;
            float yaw=-box.z_rotation;

            std::vector<Point3D> vertices = calculateVertices(x, y, z, w, l, h, yaw);
            std::vector<cv::Point2f> points;
            for (const auto& vertex : vertices) 
            {
                cv::Point2f projectedPoint = projectPoint(vertex, intrinsic, rotation, translation);
                // points.push_back(projectedPoint);
                // // 检查投影点是否在图像范围内
                if (projectedPoint.x >= 0 && projectedPoint.x < test_image.cols &&
                    projectedPoint.y >= 0 && projectedPoint.y < test_image.rows) {
                    // 在图像上绘制投影点
                    points.push_back(projectedPoint);
                }
            }

            std::vector<std::vector<int>> lines={{0,1},{1,2},{2,3},{3,0},{4,5},{5,6},{6,7},{7,4},{0,4},{1,5},{2,6},{3,7}};
            if (points.size()==8)
            {
                for (const auto& line : lines) 
                {
                        int startIndex = line[0];
                        int endIndex = line[1];
                        cv::line(test_image, points[startIndex], points[endIndex], cv::Scalar(0, 255, 0), 2);

                }
            }
            
        }
    }
    
}

std::vector<visualization_msgs::Marker> Fastbevdetect::compute_points(std::vector<fastbev::post::transbbox::BoundingBox> bboxes)
{
    // 计算8个顶点坐标，可视化
    std::vector<visualization_msgs::Marker>  markers;
    for (int i=0;i<bboxes.size();i++)
    {
        fastbev::post::transbbox::BoundingBox box=bboxes[i];
        float x=box.position.x;
        float y=box.position.y;
        float z=box.position.z;
        z=z+1;
        float w=box.size.w;
        float l=box.size.l;
        float h=box.size.h;
        float yaw=-box.z_rotation;
        

        std::cout << " x " << x 
        << " y " << y 
        << " z " << z 
        << " w " << w 
        << " l " << l 
        << " h " << h 
        << " yaw " << yaw 
        << " score " << box.score <<std::endl;

        std::vector<Point3D> vertices = calculateVertices(x, y, z, w, l, h, yaw);

        visualization_msgs::Marker marker;
        marker.header.frame_id = visual_frame;
        marker.lifetime = ros::Duration(0.5); 
        // marker.header.stamp = ros::Time::now();
        marker.ns = "object";
        marker.id = i;
        marker.type = visualization_msgs::Marker::LINE_LIST;

        marker.scale.x = 0.1;
        marker.scale.y = 0.1;
        marker.scale.z = 0.1;
        marker.color.r = 0.0;
        marker.color.g = 1.0;
        marker.color.b = 0.0;
        marker.color.a = 1.0;


        geometry_msgs::Point p1;
        p1.x=vertices[0].x;
        p1.y=vertices[0].y;
        p1.z=vertices[0].z;

        geometry_msgs::Point p2;
        p2.x=vertices[1].x;
        p2.y=vertices[1].y;
        p2.z=vertices[1].z;

        geometry_msgs::Point p3;
        p3.x=vertices[2].x;
        p3.y=vertices[2].y;
        p3.z=vertices[2].z;

        geometry_msgs::Point p4;
        p4.x=vertices[3].x;
        p4.y=vertices[3].y;
        p4.z=vertices[3].z;

        geometry_msgs::Point p5;
        p5.x=vertices[4].x;
        p5.y=vertices[4].y;
        p5.z=vertices[4].z;

        geometry_msgs::Point p6;
        p6.x=vertices[5].x;
        p6.y=vertices[5].y;
        p6.z=vertices[5].z;

        geometry_msgs::Point p7;
        p7.x=vertices[6].x;
        p7.y=vertices[6].y;
        p7.z=vertices[6].z;

        geometry_msgs::Point p8;
        p8.x=vertices[7].x;
        p8.y=vertices[7].y;
        p8.z=vertices[7].z;

        // 两个点一条线
        marker.points.push_back(p1);
        marker.points.push_back(p2);
        marker.points.push_back(p2);
        marker.points.push_back(p3);
        marker.points.push_back(p3);
        marker.points.push_back(p4);
        marker.points.push_back(p4);
        marker.points.push_back(p1);

        marker.points.push_back(p5);
        marker.points.push_back(p6);
        marker.points.push_back(p6);
        marker.points.push_back(p7);
        marker.points.push_back(p7);
        marker.points.push_back(p8);
        marker.points.push_back(p8);
        marker.points.push_back(p5);

        marker.points.push_back(p1);
        marker.points.push_back(p5);
        marker.points.push_back(p2);
        marker.points.push_back(p6);
        marker.points.push_back(p3);
        marker.points.push_back(p7);
        marker.points.push_back(p4);
        marker.points.push_back(p8);

        marker.points.push_back(p1);
        marker.points.push_back(p6);
        marker.points.push_back(p2);
        marker.points.push_back(p5);


        markers.push_back(marker);
    }
    return markers;
}


std::vector<Point3D> Fastbevdetect::calculateVertices(float x, float y, float z, float w, float l, float h, float yaw) {
    // 计算顶点的函数
    std::vector<Point3D> vertices(8);

    float half_w = w / 2.0;
    float half_l = l / 2.0;
    float half_h = h / 2.0;

    // 未旋转的顶点
    vertices[0] = {x - half_w, y - half_l, z - half_h};
    vertices[1] = {x + half_w, y - half_l, z - half_h};
    vertices[2] = {x + half_w, y + half_l, z - half_h};
    vertices[3] = {x - half_w, y + half_l, z - half_h};
    vertices[4] = {x - half_w, y - half_l, z + half_h};
    vertices[5] = {x + half_w, y - half_l, z + half_h};
    vertices[6] = {x + half_w, y + half_l, z + half_h};
    vertices[7] = {x - half_w, y + half_l, z + half_h};

    // 旋转矩阵
    double cos_yaw = std::cos(yaw);
    double sin_yaw = std::sin(yaw);
    double rotation_matrix[3][3] = {
        {cos_yaw, -sin_yaw, 0},
        {sin_yaw, cos_yaw, 0},
        {0, 0, 1}
    };


    // 应用旋转矩阵
    for (auto& vertex : vertices) {
        double x_offset = vertex.x - x;
        double y_offset = vertex.y - y;
        double z_offset = vertex.z - z;

        vertex.x = x + (cos_yaw * x_offset - sin_yaw * y_offset);
        vertex.y = y + (sin_yaw * x_offset + cos_yaw * y_offset);
        vertex.z = z + z_offset;  // z坐标保持不变
    }

    return vertices;
}


std::shared_ptr<fastbev::Core> Fastbevdetect::create_core(const std::string& model, const std::string& precision) {
  printf("Create by %s, %s\n", model.c_str(), precision.c_str());
  fastbev::pre::NormalizationParameter normalization;
  normalization.image_width = 1920;
  normalization.image_height = 1080;
  normalization.output_width = 704;
  normalization.output_height = 256;
  normalization.num_camera = 6;
  normalization.resize_lim = 0.44f;
  normalization.interpolation = fastbev::pre::Interpolation::Nearest;

  float mean[3] = {123.675, 116.28, 103.53};
  float std[3] = {58.395, 57.12, 57.375};
  normalization.method = fastbev::pre::NormMethod::mean_std(mean, std, 1.0f, 0.0f);
  fastbev::pre::GeometryParameter geo_param;
  geo_param.feat_height = 64;
  geo_param.feat_width = 176;
  geo_param.num_camera =6;
  geo_param.valid_points = 160000;
  geo_param.volum_x = 200;
  geo_param.volum_y = 200;
  geo_param.volum_z = 4;

  fastbev::CoreParameter param;
  param.pre_model = nv::format("%s/build/fastbev_pre_trt_ptq.plan", model.c_str());
  param.normalize = normalization;
  param.post_model = nv::format("%s/build/fastbev_post_trt_ptq.plan", model.c_str());
  param.geo_param = geo_param;
  return fastbev::create_core(param);
}